import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset

N = 1000
batch_size = 32
seq_len = 5  # 每条样本包含 5 个时间步
input_size = 5  # 每个时间步是 5 维向量
hidden_size = 8

# 输入维度为5维向量,输出维度为8维向量
cell = nn.RNNCell(input_size=input_size, hidden_size=hidden_size)

# 生成1000个样本，每个样本有5个时间步，每个时间步是5维向量
X = torch.randn(N, seq_len, input_size)
print(X.shape)

# 将1000个样本放入DataLoader中进行批量处理,每个批次的样本数量为32,并且打乱顺序
loader = DataLoader(TensorDataset(X), batch_size=batch_size, shuffle=True)

for i, (x,) in enumerate(loader):
    B = x.size(0)
    h = torch.zeros(B, hidden_size)
    # x 的形状是 [32, 5, 5]，代表32个样本，每个样本有5个时间步，每个时间步是5维向量
    # x_t = x[:, t, :] 通过切片操作提取所有32个样本在时间步 t 的数据
    # x_t 的形状变为 [32, 5]，即32个样本在同一时间步的输入
    # cell(x_t, h) 同时对这32个样本进行RNN计算
    for t in range(seq_len):
        x_t = x[:, t, :]
        h = cell(x_t, h)
    print(f"[RNNCell] batch {i}: xb={x.shape} -> h_T={h.shape}")
    break
